#!/usr/bin/env python
# coding=UTF-8
'''
@Author: Jake Gu
@Date: 2019-04-29 20:56:42
@LastEditTime: 2019-04-29 21:01:34
'''
##
# @file   setup.py.in
# @author Yibo Lin
# @date   Jun 2018
# @brief  For CMake to generate setup.py file 
#

import os 
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension

def add_prefix(filename):
    return os.path.join('${CMAKE_CURRENT_SOURCE_DIR}/src', filename)

modules = []

modules.extend([
    CppExtension('dct_cpp', 
        [
            add_prefix('dct.cpp'),
            add_prefix('dst.cpp'),
            add_prefix('dxt.cpp'), 
            add_prefix('dct_2N.cpp')
            ]), 
    CppExtension('dct_lee_cpp', 
        [
            add_prefix('dct_lee.cpp')
            ]),
            ])

if not "${CUDA_FOUND}" or "${CUDA_FOUND}".upper() == 'TRUE': 
    modules.extend([
            CUDAExtension('dct_cuda', 
                [
                    add_prefix('dct_cuda.cpp'),
                    add_prefix('dct_cuda_kernel.cu'),
                    add_prefix('dst_cuda.cpp'),
                    add_prefix('dst_cuda_kernel.cu'),
                    add_prefix('dxt_cuda.cpp'), 
                    add_prefix('dct_2N_cuda.cpp')
                    ]),
            CUDAExtension('dct_lee_cuda', 
                [
                    add_prefix('dct_lee_cuda.cpp'),
                    add_prefix('dct_lee_cuda_kernel.cu'), 
                    add_prefix('dct_cuda_kernel.cu'),
                    add_prefix('dst_cuda_kernel.cu')
                    ]),
            CUDAExtension('dct2_fft2d_cuda', 
                [
                    add_prefix('dct2_fft2d_cuda.cpp'),
                    add_prefix('dct2_fft2d_cuda_kernel.cu')
                    ]),
        ])

setup(
        name='dct',
        ext_modules=modules,
        cmdclass={
            'build_ext': BuildExtension
            })
